import os
import mysql.connector
from mysql.connector import Error
from mysql.connector import pooling
from typing import List, Dict, Any, Optional, Union, ContextManager
import threading
import time
import logging
from Config import *

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('MySQLPool')


class MySQLPool:
    """MySQL 数据库连接池实现"""

    def __init__(self,
                 host: str,
                 user: str,
                 password: str,
                 database: str,
                 port: Optional[int] = None,  # 可选端口参数
                 pool_name: str = "mysql_pool",
                 pool_size: int = 5,
                 max_overflow: int = 5,
                 timeout: float = 30.0):
        """
        初始化数据库连接池

        Args:
            host: 数据库主机地址
            user: 数据库用户名
            password: 数据库密码
            database: 数据库名
            port: 数据库端口，默认从环境变量获取或使用MySQL默认端口(3306)
            pool_name: 连接池名称
            pool_size: 连接池初始大小
            max_overflow: 允许的最大溢出连接数
            timeout: 获取连接的超时时间（秒）
        """
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.port = self._get_port(port)
        self.pool_name = pool_name
        self.pool_size = pool_size
        self.max_overflow = max_overflow
        self.timeout = timeout

        self._pool = None
        self._overflow = 0
        self._lock = threading.Lock()
        self._init_pool()

    def _get_port(self, port: Optional[int]) -> int:
        """获取数据库端口，优先使用传入参数，其次是环境变量，最后是默认值"""
        if port is not None:
            return port

        env_port = os.getenv('MYSQL_PORT')
        if env_port:
            try:
                return int(env_port)
            except ValueError:
                logger.warning(f"无效的端口环境变量: {env_port}，使用默认端口 3306")

        return 3306  # MySQL默认端口

    def _init_pool(self) -> None:
        """初始化连接池"""
        try:
            self._pool = pooling.MySQLConnectionPool(
                pool_name=self.pool_name,
                pool_size=self.pool_size,
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database,
                port=self.port
            )
            logger.info(f"连接池 '{self.pool_name}' 初始化成功，端口: {self.port}，大小: {self.pool_size}")
        except Error as e:
            logger.error(f"连接池初始化失败: {e}")
            raise

    def get_connection(self) -> 'MySQLConnectionWrapper':
        """获取数据库连接"""
        start_time = time.time()

        while True:
            with self._lock:
                # 尝试从池中获取连接
                try:
                    connection = self._pool.get_connection()
                    if connection.is_connected():
                        logger.debug("从连接池获取连接成功")
                        return MySQLConnectionWrapper(connection, self)
                except Error as e:
                    logger.warning(f"从连接池获取连接失败: {e}")

                # 如果池已满且未达到最大溢出，创建溢出连接
                if self._overflow < self.max_overflow:
                    try:
                        connection = mysql.connector.connect(
                            host=self.host,
                            user=self.user,
                            password=self.password,
                            database=self.database
                        )
                        self._overflow += 1
                        logger.info(f"创建溢出连接，当前溢出: {self._overflow}/{self.max_overflow}")
                        return MySQLConnectionWrapper(connection, self, is_overflow=True)
                    except Error as e:
                        logger.error(f"创建溢出连接失败: {e}")

            # 等待并重试
            elapsed = time.time() - start_time
            if elapsed >= self.timeout:
                raise TimeoutError(f"获取数据库连接超时 ({self.timeout}s)")

            wait_time = min(1.0, self.timeout - elapsed)
            time.sleep(wait_time)

    def release(self, connection: mysql.connector.MySQLConnection, is_overflow: bool = False) -> None:
        """释放数据库连接"""
        with self._lock:
            try:
                if connection.is_connected():
                    if is_overflow:
                        connection.close()
                        self._overflow -= 1
                        logger.info(f"关闭溢出连接，当前溢出: {self._overflow}/{self.max_overflow}")
                    else:
                        connection.close()  # 返回到池中
                        logger.debug("连接已返回连接池")
                else:
                    logger.warning("尝试释放已关闭的连接")
            except Error as e:
                logger.error(f"释放连接失败: {e}")

    def close(self) -> None:
        """关闭连接池"""
        with self._lock:
            try:
                # 关闭所有溢出连接
                while self._overflow > 0:
                    logger.info(f"清理溢出连接，剩余: {self._overflow}")
                    self._overflow = 0

                # 关闭连接池
                if self._pool:
                    self._pool._remove_connections()
                    self._pool = None
                    logger.info(f"连接池 '{self.pool_name}' 已关闭")
            except Error as e:
                logger.error(f"关闭连接池失败: {e}")


class MySQLConnectionWrapper:
    """数据库连接包装器，实现上下文管理器"""

    def __init__(self, connection: mysql.connector.MySQLConnection, pool: MySQLPool, is_overflow: bool = False):
        self._connection = connection
        self._pool = pool
        self._is_overflow = is_overflow
        self._cursor = None

    def __enter__(self) -> 'MySQLConnectionWrapper':
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    def execute(self, query: str, params: Optional[Dict[str, Any]] = None) -> int:
        """执行SQL语句（INSERT, UPDATE, DELETE）"""
        try:
            self._ensure_cursor()
            self._cursor.execute(query, params or {})
            self._connection.commit()
            return self._cursor.rowcount
        except Error as e:
            self._connection.rollback()
            logger.error(f"执行SQL失败: {e}")
            raise

    def query(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        """执行查询语句（SELECT）"""
        try:
            self._ensure_cursor(dictionary=True)
            self._cursor.execute(query, params or {})
            return self._cursor.fetchall()
        except Error as e:
            logger.error(f"查询失败: {e}")
            raise

    def _ensure_cursor(self, dictionary: bool = False) -> None:
        """确保游标已创建"""
        if not self._cursor or self._cursor._closed:
            self._cursor = self._connection.cursor(dictionary=dictionary)

    def close(self) -> None:
        """关闭连接（返回池中）"""
        if self._cursor and not self._cursor._closed:
            self._cursor.close()

        self._pool.release(self._connection, self._is_overflow)
        self._connection = None
        self._cursor = None

dbpool = MySQLPool(
        DATASOURCE_HOST,
        DATASOURCE_USERNAME,
        DATASOURCE_PASSWORD,
        DATASOURCE_DATABASE,
        DATASOURCE_PORT
    )

# 使用示例
# if __name__ == "__main__":
#     # 配置数据库连接参数
#     pool = MySQLPool(
#         host="localhost",
#         user="your_username",
#         password="your_password",
#         database="your_database",
#         pool_size=3,
#         max_overflow=2
#     )
#
#     # 使用上下文管理器获取连接
#     with pool.get_connection() as conn:
#         # 执行查询
#         results = conn.query("SELECT * FROM users LIMIT 5")
#         for row in results:
#             print(row)
#
#         # 执行插入
#         conn.execute("INSERT INTO users (name, email) VALUES (%(name)s, %(email)s)",
#                      {"name": "Test User", "email": "test@example.com"})
#
#     # 模拟多线程使用连接池
#     import threading
#
#
#     def worker():
#         with pool.get_connection() as conn:
#             result = conn.query("SELECT COUNT(*) as count FROM users")
#             print(f"线程 {threading.get_ident()}: 用户数量 = {result[0]['count']}")
#
#
#     threads = []
#     for _ in range(5):
#         t = threading.Thread(target=worker)
#         threads.append(t)
#         t.start()
#
#     for t in threads:
#         t.join()
#
#     # 关闭连接池
#     pool.close()